import torch
from torch import Tensor
import torch.nn as nn
import os


class Model(nn.Module):
    def __init__(self, net: nn.Module = None, optimizer: torch.optim.Optimizer = None, criterion=None) -> None:
        """
        Model 类，保存着网络、优化器、损失函数，为了尽可能的提高代码的可读性，这里将这三个东西都放在了一起，方便查看；
        为了提高代码的自由度，这里没有将这三个东西写成类的属性，而是写成了类的参数，这样可以在实例化的时候传入不同的网络、优化器、损失函数
        :param net: nn.Module 神经网络
        :param optimizer: torch.optim.Optimizer 优化器
        :param criterion: 损失函数
        """
        super().__init__()
        self.net = net
        self.optimizer = optimizer
        self.criterion = criterion

    def forward(self, state: Tensor) -> Tensor:
        """
        继承自 nn.Module 的 forward 函数，用于前向传播
        :param state: Tensor 状态
        :return: Tensor 返回模型的输出
        """
        return self.net(state)

    def update(self, loss: Tensor) -> None:
        """
        更新网络
        :param loss: Tensor 损失，这里的损失是在外部计算好的，可以利用model.criterion()函数计算损失
        :return: None
        """
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

    def save(self, path: str) -> None:
        """
        保存网络
        :param path: str 保存的文件名
        :return: None
        """
        dir_name = os.path.dirname(path)
        if not os.path.exists(dir_name):
            os.mkdir(dir_name)
        torch.save(self.net.state_dict(), path)

    def load(self, path: str) -> bool:
        """
        加载网络
        :param path: str 保存的文件名
        :return: bool 是否加载成功
        """
        if os.path.exists(path):
            self.net.load_state_dict(torch.load(path))
            return True
        else:
            return False
